package com.agent.rag.config;

import com.agent.commons.errorcode.MilvusErrorCodeConstants;
import com.agent.commons.exception.BaseException;
import io.milvus.common.clientenum.FunctionType;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.DataType;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.collection.request.AddFieldReq;
import io.milvus.v2.service.collection.request.CreateCollectionReq;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import static com.agent.commons.constants.MilvusConstants.*;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * Milvus集合及架构配置
 */
@Slf4j
public class MilvusCollectionConfig {
    @Resource
    private MilvusClientV2 milvus;


    /**
     * 构建分词器(采用结巴分词)
     * @return
     */
    public Map<String,Object> buildAnalyzer(){
        Map<String,Object> analyzer=new HashMap<>();
        analyzer.put("tokenizer", "jieba");
        return analyzer;
    }

    /**
     * 创建Milvus集合架构
     * @param analyzer
     * @return
     */
    @Bean
    public CreateCollectionReq.CollectionSchema createMilvusSchema(Map<String,Object> analyzer){
        CreateCollectionReq.CollectionSchema schema = milvus.createSchema();
        //添加主键id
        schema.addField(AddFieldReq.builder()
                .fieldName("id")
                .dataType(DataType.Int64)//设置主键数据类型
                .isPrimaryKey(true)//设置主键
                .autoID(true)//自增策略
                .build());
        //添加用户id字段
        schema.addField(AddFieldReq.builder()
                .fieldName("userId")
                .isPartitionKey(true)//设置为分区键
                .dataType(DataType.VarChar)
                .build());
        //添加文本字段
        schema.addField(AddFieldReq.builder()
                .fieldName("text")
                .dataType(DataType.VarChar)
                .enableAnalyzer(true)//启用分词器
                .analyzerParams(analyzer)//设置分词器
                .build());
        //添加文档图片包含
        schema.addField(AddFieldReq.builder()
                .fieldName("photo")
                .dataType(DataType.Array)
                .build());
        //添加文本密集矢量化字段
        schema.addField(AddFieldReq.builder()
                .fieldName("dense_vector")
                .dataType(DataType.FloatVector)
                .dimension(1024)
                .build());
        //添加文本稀疏矢量化字段
        schema.addField(AddFieldReq.builder()
                .fieldName("sparse_vector")
                .dataType(DataType.SparseFloatVector)
                .dimension(1024)
                .build());
        //添加BM25稀疏矢量方法
        schema.addFunction(CreateCollectionReq.Function.builder()
                .functionType(FunctionType.BM25)
                .name("bm25")
                .inputFieldNames(List.of("text"))
                .outputFieldNames(List.of("sparse_vector"))
                .build());
        return schema;
    }

    /**
     * Milvus字段索引创建
     */
    public List<IndexParam> createMilvusIndex(){
        log.info("Milvus集合索引创建成功！");
        IndexParam index_one=IndexParam.builder()
                .fieldName("dense_vector")
                .indexName("dense")
                .indexType(IndexParam.IndexType.IVF_FLAT)//采用倒排索引
                .metricType(IndexParam.MetricType.IP)//采用内积计算相似度
                .build();
        IndexParam index_two=IndexParam.builder()
                .fieldName("sparse_vector")
                .indexName("sparse")
                .indexType(IndexParam.IndexType.SPARSE_INVERTED_INDEX)//采用稀疏向量索引
                .metricType(IndexParam.MetricType.BM25)//采用BM25算法计算相似度
                .build();
        return List.of(index_one,index_two);
    }

    /**
     * 创建Milvus集合
     */
    @Bean
    public boolean createCollection(List<IndexParam> index,Map<String,Object> analyzer){
        try {
            milvus.useDatabase(DATABASE_NAME);
        } catch (InterruptedException e) {
            throw new BaseException(MilvusErrorCodeConstants.DATABASE_CONNECTION_ERROR);
        }
        //创建集合
        milvus.createCollection(CreateCollectionReq.builder()
                .collectionName(TEXT_COLLECTION)
                //.collectionSchema(schema)//指定集合名称
                .indexParams(index)//创建集合索引
                .numPartitions(64)//指定分区数
                //.properties(properties)//指定分区键隔离机制
                .build());
        log.info("Milvus集合创建成功！");
        return true;
    }
}
